package demo.netflix.discovery.util;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 基于**Token Bucket Algorithm ( 令牌桶算法 )**的速率限制器。
 * <p>
 * 令牌桶算法的原理是系统会以一个恒定的速度往桶里放入令牌，而如果请求需要被处理，则需要先从桶里获取一个令牌，当桶里没有令牌可取时，则拒绝服务。
 * 
 * Rate limiter implementation is based on token bucket algorithm. There are two
 * parameters:
 * <ul>
 * <li>burst size - maximum number of requests allowed into the system as a
 * burst</li>
 * <li>average rate - expected number of requests per second (RateLimiters using
 * MINUTES is also supported)</li>
 * </ul>
 * 举个 例子: 来理解这两个参数 + 构造方法里的一个参数：
 * <li>averageRateUnit = SECONDS
 * <li>averageRate = 2000 ,每 秒 可获取 2000 个令牌。例如，每秒允许请求 2000 次。每 毫秒 可填充 2000 /
 * 1000 = 2 个 消耗 的令牌。
 * <li>burstSize = 10,每 毫秒 可获取 10 个令牌。
 * <p>
 * 例如，每毫秒允许请求上限为 10 次，并且请求消耗掉的令牌，需要逐步填充。 这里要注意下，虽然每毫秒允许请求上限为 10 次，这是在没有任何令牌被 消耗
 * 的情况下，实际每秒允许请求依然是 2000 次。 这就是基于令牌桶算法的限流的特点：让流量平稳，而不是瞬间流量。1000 QPS
 * 相对平均的分摊在这一秒内，而不是第 1 ms 999 请求，后面 999 ms 0 请求 。
 */
public class RateLimiter {

    /**
     * 速率单位转换成毫秒
     */
    private final long rateToMsConversion;

    /**
     * 消耗令牌数
     */
    private final AtomicInteger consumedTokens = new AtomicInteger();

    /**
     * 最后填充令牌的时间
     */
    private final AtomicLong lastRefillTime = new AtomicLong(0);

    /**
     * 创建一个限速器
     * 
     * @param averageRateUnit 速率单位，目前支持 分钟级 和 秒级 两种速率限制。
     */
    public RateLimiter(TimeUnit averageRateUnit) {
        switch (averageRateUnit) {
        case SECONDS:
            rateToMsConversion = 1000;
            break;
        case MINUTES:
            rateToMsConversion = 60 * 1000;
            break;
        default:
            throw new IllegalArgumentException("TimeUnit of " + averageRateUnit + " is not supported");
        }
    }

    /**
     * 获取令牌
     * 
     * @param burstSize   令牌桶上限
     * @param averageRate 令牌再装平均速率
     * @return 是否获取成功
     */
    public boolean acquire(int burstSize, long averageRate) {
        return acquire(burstSize, averageRate, System.currentTimeMillis());
    }

    /**
     * 获取令牌
     * 
     * @param burstSize         令牌桶上限
     * @param averageRate       令牌再装平均速率
     * @param currentTimeMillis 当前时间毫秒
     * @return 是否获取成功
     */
    public boolean acquire(int burstSize, long averageRate, long currentTimeMillis) {
        // Instead of throwing exception, we just let all the traffic go
        if (burstSize <= 0 || averageRate <= 0) {
            return true;
        }
        // 填充 已消耗 的令牌。
        refillToken(burstSize, averageRate, currentTimeMillis);
        return consumeToken(burstSize);
    }

    /**
     * 填充 令牌 。
     * <p>
     * 为什么不起一个后台线程执行填充呢？
     * <li>一方面，实际项目里每个接口都会有相应的 RateLimiter ，导致 太多 执行频率 极高 的后台任务；
     * <li>另一方面，获取令牌时才计算，多次令牌填充可以合并成一次，减少冗余和无效的计算。
     * 
     * @param burstSize         令牌桶上限
     * @param averageRate       令牌填充 平均 速率
     * @param currentTimeMillis
     */
    private void refillToken(int burstSize, long averageRate, long currentTimeMillis) {
        // 获得 最后填充令牌的时间
        long refillTime = lastRefillTime.get();
        // 获得 过去多少毫秒，用于计算需要填充的令牌数
        long timeDelta = currentTimeMillis - refillTime;
        // 计算 可填充最大令牌数量，可能超过 burstSize ，所以下面会有逻辑调整 newTokens
        long newTokens = timeDelta * averageRate / rateToMsConversion;
        if (newTokens > 0) {
            /*
             * 计算 新的填充令牌的时间。为什么不能用 currentTimeMillis 呢 ？ 例如， averageRate = 500 &&
             * averageRateUnit = SECONDS 时， 每 2 毫秒才填充一个令牌，如果设置 currentTimeMillis ，
             * 会导致不足以填充一个令牌的时长被吞了 。
             */
            long newRefillTime = refillTime == 0 ? currentTimeMillis
                    : refillTime + newTokens * rateToMsConversion / averageRate;
            // CAS 保证有且仅有一个线程进入填充
            if (lastRefillTime.compareAndSet(refillTime, newRefillTime)) {
                while (true) {// 死循环，直到成功
                    // 计算 填充令牌后的已消耗令牌数量
                    int currentLevel = consumedTokens.get();
                    // burstSize 可能调小，例如，系统接入分布式配置中心，可以远程调整该数值。如果此时 burstSize 更小，以它作为 已消耗 的令牌数量。
                    int adjustedLevel = Math.min(currentLevel, burstSize); // In case burstSize decreased
                    int newLevel = (int) Math.max(0, adjustedLevel - newTokens);
                    // CAS 避免和正在消费令牌的线程冲突
                    if (consumedTokens.compareAndSet(currentLevel, newLevel)) {
                        return;
                    }
                }
            }
        }
    }

    /**
     * 消费 令牌
     * 
     * @param burstSize 令牌桶上限。
     * @return
     */
    private boolean consumeToken(int burstSize) {
        // 死循环，直到没有令牌，或者获取令牌成功
        while (true) {
            // 没有令牌
            int currentLevel = consumedTokens.get();
            if (currentLevel >= burstSize) {
                return false;
            }
            // CAS 避免和正在消费令牌或者填充令牌的线程冲突
            if (consumedTokens.compareAndSet(currentLevel, currentLevel + 1)) {
                return true;
            }
        }
    }

    public void reset() {
        consumedTokens.set(0);
        lastRefillTime.set(0);
    }
}
